/// 一个简易 LALR(1) 文法解析器
/// 输入：文法 G
/// 输出：ACTION，GOTO 表
/// 输入样例：
/// 第一行两个非负整数 n,m, 接下来有 1 个空串表示符，1 个文法开始符，有 n 个终端符，m 个产生式
/// 非终端符和终端符之间以 whitespace 区分，目前不支持 | 运算，例如下
///
/// 12 10
/// nul <program>
/// <data_type>
/// <id>
/// <int>
/// <unary_op>
/// "("
/// ")"
/// "{"
/// "}"
/// ";"
/// "return"
/// ("+"|"-")
/// ("*"|"/")
/// <program> ::= <function>
/// <function> ::= <data_type> <id> "(" ")" "{" <statement> "}"
/// <statement> ::= "return" <exp> ";"
/// <exp> ::= <term>
/// <exp> ::= <exp> ("+"|"-") <term>
/// <term> ::= <factor>
/// <term> ::= <term> ("*"|"/") <factor>
/// <factor> ::= <int>
/// <factor> ::= <unary_op> <factor>
/// <factor> ::= "(" <exp> ")"
use std::{
    collections::{BTreeMap, BTreeSet, HashMap, VecDeque},
    fs::File,
    io::{BufRead, BufReader},
    usize, vec,
};

/// 在向前看符号中，特别地用 usize::MAX 表示 EOF，EOF 不会加入实际的符号表
#[derive(Debug)]
pub struct Grammer {
    sym: Vec<String>,
    prod: Vec<Production>,

    // 以下是辅助的 ds
    sym_map: BTreeMap<String, usize>,
    term_vis: Vec<bool>,
    nul: usize,
    start: usize,
    related_lhs_prod_id: Vec<BTreeSet<usize>>,

    first: Vec<BTreeSet<usize>>,
    goto: HashMap<(Vec<(usize, usize, usize)>, usize), Vec<(usize, usize, usize)>>,
    lr0_kernels: BTreeSet<Vec<(usize, usize)>>,
    lalr1_all_lookahead: Vec<BTreeSet<usize>>,
    lalr1_lookahead_set_to_id: BTreeMap<BTreeSet<usize>, usize>,
    lalr1_transmit_map: BTreeMap<Vec<(usize, usize)>, Vec<Vec<(usize, usize)>>>,
    lalr1_lookahead_map: BTreeMap<Vec<(usize, usize)>, usize>,
}

#[derive(Debug)]
pub struct Production {
    pub ret: usize,
    pub rhs: Vec<usize>,
}

impl Grammer {
    /// 获取符号 s 是否是终结符
    pub fn is_term(&self, s: usize) -> bool {
        self.term_vis[s]
    }

    /// 获取符号 s 是否为空表示符
    pub fn is_null(&self, s: usize) -> bool {
        s == self.nul
    }

    /// 获取符号 s 是否为 EOF
    pub fn is_eof(&self, s: usize) -> bool {
        s == usize::MAX
    }

    /// 获取不同符号的个数
    pub fn sym_count(&self) -> usize {
        self.sym.len()
    }

    /// 获取文法中产生式的个数
    pub fn prod_count(&self) -> usize {
        self.prod.len()
    }

    pub fn get_eof(&self) -> usize {
        usize::MAX
    }

    /// 获取全体符号的 First 集
    pub fn get_first_all(&mut self) {
        self.first.resize(self.sym_count(), BTreeSet::new());

        for s in 0..self.sym_count() {
            if self.is_term(s) || self.is_null(s) {
                self.first[s].insert(s);
            }
        }

        loop {
            let mut ok = true;
            for p in 0..self.prod_count() {
                let pr = &self.prod[p];
                let ret = self.get_first(&pr.rhs);
                let psize = self.first[pr.ret].len();
                self.first[pr.ret].extend(ret);
                if self.first[pr.ret].len() > psize {
                    ok = false;
                }
            }

            if ok {
                break;
            }
        }
    }

    /// 获取序列 seq 的 First 集
    pub fn get_first(&self, seq: &Vec<usize>) -> BTreeSet<usize> {
        let mut ret: BTreeSet<usize> = BTreeSet::new();
        let mut all_null = true;

        // 找到第一个推导非空的符号
        for sym in seq {
            let r0 = self.first[*sym].clone();
            ret.extend(r0.iter().filter(|&&s| !self.is_null(s)));

            if !r0.contains(&self.nul) {
                all_null = false;
                break;
            }
        }

        // 能推导全空，加空进去
        if all_null {
            ret.insert(self.nul);
        }

        ret
    }

    /// 构建项目集 i 的 LR(0) 闭包
    /// 在 LR(1) 中，一个项目集形如 8 <factor> ::= <unary_op> . <factor> , $ 可以由三个整数来表示
    /// (产生式编号，光标位置，向前看符号), 上面这个就应该是 (8,1,$)
    /// 这里采用优化的方法，先计算的是 LR(0) 闭包，没有向前看
    /// 统一为 (8,1,0) 即可
    pub fn get_lr0_closure(&self, mut i: Vec<(usize, usize)>) -> Option<Vec<(usize, usize)>> {
        let mut res: Vec<(usize, usize)> = Vec::new();
        let mut vis: Vec<bool> = Vec::new();
        vis.resize(self.prod.len(), false);
        res.append(&mut i);

        // 一直展开对应项目集位置非终端符，直到没有项目集可加入为止
        loop {
            let mut added: Vec<(usize, usize)> = vec![];

            for (prod_id, cur) in &res {
                // 越界
                if *prod_id >= self.prod.len() {
                    return None;
                }

                let rhs = &self.prod[*prod_id].rhs;
                if *cur > rhs.len() {
                    return None;
                }

                // 看起来是一个归约状态
                if *cur == rhs.len() {
                    continue;
                }

                // 遍历以符号 rhs[*cur] 作为产生式 lhs 的所有产生式
                for sub_prod in &self.related_lhs_prod_id[rhs[*cur]] {
                    if vis[*sub_prod] {
                        continue;
                    }

                    vis[*sub_prod] = true;
                    added.push((*sub_prod, 0));
                }
            }

            if added.is_empty() {
                break;
            }
            res.append(&mut added);
        }

        Some(res)
    }

    /// LR(0) GOTO 函数
    /// 输入项目集族与文法符号，给出 GOTO(item,X) 的结果
    pub fn get_lr0_goto(
        &self,
        item: &Vec<(usize, usize)>,
        x: usize,
    ) -> Option<Vec<(usize, usize)>> {
        let mut res: Vec<(usize, usize)> = Vec::new();
        for (prod_id, cur) in item {
            if let Some(nx) = self.get_next_symbol((*prod_id, *cur)) {
                if x == nx {
                    res.push((*prod_id, *cur + 1));
                }
            } else {
                return None;
            }
        }
        let ret = self.get_lr0_closure(res);

        ret
    }

    /// 获取当前项目的下一个符号
    pub fn get_next_symbol(&self, now: (usize, usize)) -> Option<usize> {
        let (prod_id, cur) = now;
        if prod_id >= self.prod.len() {
            return None;
        }
        // 归约状态下没有下一个符号
        if cur >= self.prod[prod_id].rhs.len() {
            return None;
        }

        Some(self.prod[prod_id].rhs[cur])
    }

    /// 构建 lr0 内核
    pub fn build_lr0_kernels(&mut self) {
        self.lr0_kernels.clear();
        let init = vec![(0, 0)];
        let mut que: VecDeque<Vec<(usize, usize)>> = VecDeque::new();
        que.push_back(init);

        // 构建完整 lr0 项目集族
        while let Some(item) = que.pop_front() {
            if let Some(close) = self.get_lr0_closure(item) {
                self.lr0_kernels.insert(close.clone());
                for x in 0..self.sym.len() {
                    if let Some(nw) = self.get_lr0_goto(&close, x) {
                        if self.lr0_kernels.get(&nw).is_none() {
                            self.lr0_kernels.insert(nw.clone());
                            que.push_back(nw);
                        }
                    }
                }
            }
        }

        // 筛选内核项
        let mut filtered_kernels = BTreeSet::new();
        for kernel in self.lr0_kernels.iter() {
            let filtered = self.lr0_item_to_kernel(kernel);

            // 保留非空内核
            if !filtered.is_empty() {
                filtered_kernels.insert(filtered);
            }
        }

        self.lr0_kernels = filtered_kernels;
    }

    /// 把 lr0 项目压缩为 lr0 内核
    pub fn lr0_item_to_kernel(&self, item: &Vec<(usize, usize)>) -> Vec<(usize, usize)> {
        let mut filtered: BTreeSet<(usize, usize)> = BTreeSet::new();

        for &(prod_idx, dot_pos) in item {
            if self.is_lr0_kernel(prod_idx, dot_pos) {
                filtered.insert((prod_idx, dot_pos));
            }
        }

        filtered.into_iter().collect()
    }

    pub fn is_lr0_kernel(&self, p: usize, d: usize) -> bool {
        p == 0 || d > 0
    }

    /// 构建 LALR(1) 内核项
    pub fn build_lalr1_kernels(&mut self) {
        let mut lookahead_set_to_id: BTreeMap<BTreeSet<usize>, usize> = BTreeMap::new();
        let mut all_lookahead: Vec<BTreeSet<usize>> = vec![];

        // 预处理 LR(0) 项和 FIRST 集
        self.get_first_all();
        self.build_lr0_kernels();

        // 第 0 项表示为 {}
        // 第 1 项的 LA 表示为 {$}
        let mut s_eof = BTreeSet::new();
        lookahead_set_to_id.insert(s_eof.clone(), 0);
        all_lookahead.push(s_eof.clone());
        s_eof.insert(self.get_eof());
        lookahead_set_to_id.insert(s_eof.clone(), 1);
        all_lookahead.push(s_eof);
        self.lalr1_lookahead_map.insert(vec![(0, 0)], 1);

        // 构建 LA 传播图
        for k in &self.lr0_kernels {
            // 内核转闭包
            if let Some(c) = self.get_lr0_closure(k.clone()) {
                // 对于闭包的每一个项目
                for &(p, cur) in &c {
                    // 下一个可能的文法符号
                    if let Some(nx) = self.get_next_symbol((p, cur)) {
                        // 可传播
                        if let Some(gt) = self.get_lr0_goto(&c, nx) {
                            // 项目再转内核
                            let next_ker = self.lr0_item_to_kernel(&gt);

                            // 插入到转移图中
                            self.lalr1_transmit_map
                                .entry(k.to_vec())
                                .or_insert_with(Vec::new)
                                .push(next_ker.clone());

                            // a.B\beta,a
                            // 按 First(\beta a) 产生的 LA，不过这里的 a 还没确定
                            // 确保自发产生的 LA 一定由内核项产生
                            if self.is_lr0_kernel(p, cur) && !self.is_term(nx) {
                                let rest = self.prod[p].rhs[cur + 1..].to_vec();
                                let mut fs = self.get_first(&rest);
                                fs.extend(
                                    all_lookahead[*self
                                        .lalr1_lookahead_map
                                        .entry(next_ker.clone())
                                        .or_insert(0)]
                                    .clone(),
                                );
                                fs.remove(&self.nul);

                                if !lookahead_set_to_id.contains_key(&fs) {
                                    self.lalr1_lookahead_map
                                        .insert(next_ker, lookahead_set_to_id.len());

                                    lookahead_set_to_id
                                        .insert(fs.clone(), lookahead_set_to_id.len());
                                    all_lookahead.push(fs);
                                }
                            }
                        }
                    }
                }
            }
        }

        // 传播
        let mut changed = true;
        while changed {
            changed = false;
            for (u, adj) in &self.lalr1_transmit_map {
                let u_la: BTreeSet<usize>;
                if let Some(&id) = self.lalr1_lookahead_map.get(u) {
                    u_la = all_lookahead[id].clone();
                } else {
                    // 如果传播源为空，则跳过
                    continue;
                }

                for v in adj {
                    let mut v_la = all_lookahead[*self
                        .lalr1_lookahead_map
                        .entry(v.to_vec())
                        .or_insert(lookahead_set_to_id[&BTreeSet::new()])]
                    .clone();

                    let before = v_la.len();

                    // 扩展向前看集合
                    v_la.extend(u_la.clone());

                    if v_la.len() > before {
                        changed = true;
                    }

                    // 映射集合为 usize
                    if !lookahead_set_to_id.contains_key(&v_la) {
                        lookahead_set_to_id.insert(v_la.clone(), lookahead_set_to_id.len());
                        all_lookahead.push(v_la.clone());
                    }

                    self.lalr1_lookahead_map
                        .insert(v.to_vec(), lookahead_set_to_id[&v_la]);
                }
            }
        }

        // for (idx, (a, b)) in self.lalr1_lookahead_map.iter().enumerate() {
        //     println!("LALR(1) kernel {} :", idx);
        //     let la = all_lookahead[*b].clone();

        //     for &(p, d) in a {
        //         println!("{},{},{:?}", p, d, la);
        //     }
        // }

        self.lalr1_all_lookahead = all_lookahead;
        self.lalr1_lookahead_set_to_id = lookahead_set_to_id;
    }
}

/// 针对轻量的文法输入，读入一个 Grammer
/// 此文法会自动变为扩宽文法，给定文法开始符 S
/// 则第 0 个产生式一定是 S' ::= S
/// 其中 S' 不实际出现在 symbol 中，但其映射的 usize 一定为 0
pub fn light_grammer_input(f: File) -> std::io::Result<Grammer> {
    let mut fin = BufReader::new(f);
    let mut buf = String::new();

    fin.read_line(&mut buf)?;
    let mut it1 = buf.split_ascii_whitespace();
    let n: usize = it1
        .next()
        .ok_or(std::io::ErrorKind::InvalidInput)?
        .parse()
        .or_else(|_| Err(std::io::ErrorKind::InvalidInput))?;
    let m: usize = it1
        .next()
        .ok_or(std::io::ErrorKind::InvalidInput)?
        .parse()
        .or_else(|_| Err(std::io::ErrorKind::InvalidInput))?;

    let mut map: BTreeMap<String, usize> = BTreeMap::new();

    // S'
    map.insert(String::new(), map.len());

    let mut term: Vec<String> = vec![];
    let mut prod_str: Vec<String> = vec![];

    // 先预处理符号 vec，存入 map 中再遍历转为 usize
    // 空串 文法开始符
    buf.clear();
    fin.read_line(&mut buf)?;
    it1 = buf.split_ascii_whitespace();

    let nul: String = it1
        .next()
        .ok_or(std::io::ErrorKind::InvalidInput)?
        .parse()
        .or_else(|_| Err(std::io::ErrorKind::InvalidInput))?;
    map.insert(nul.clone(), map.len());
    let start: String = it1
        .next()
        .ok_or(std::io::ErrorKind::InvalidInput)?
        .parse()
        .or_else(|_| Err(std::io::ErrorKind::InvalidInput))?;

    // 终端符
    for _ in 0..n {
        buf.clear();
        fin.read_line(&mut buf)?;
        let nw = buf.clone().trim_end().to_string();

        term.push(nw.clone());
        if !map.contains_key(&nw) {
            map.insert(nw, map.len());
        }
    }

    // 产生式
    for _ in 0..m {
        buf.clear();
        fin.read_line(&mut buf)?;
        let nw = buf.clone().trim_end().to_string();

        prod_str.push(nw);

        let it2 = buf.split_ascii_whitespace();

        for c in it2 {
            if c != "::=" {
                if !map.contains_key(c) {
                    map.insert(c.to_string(), map.len());
                }
            }
        }
    }

    let mut sym: Vec<String> = vec![];
    let mut term_vis: Vec<bool> = vec![];
    let mut ps: Vec<Production> = vec![];
    let mut related_lhs_prod_id: Vec<BTreeSet<usize>> = vec![];
    let n = map.keys().len();
    term_vis.resize(n, false);
    related_lhs_prod_id.resize(n, BTreeSet::new());

    for c in map.keys() {
        sym.push(c.to_string());
    }

    for c in term {
        // assert all the terms exist
        let v = map.get(&c).unwrap();
        term_vis[*v] = true;
    }

    ps.push(Production {
        ret: 0,
        rhs: vec![map[&start]],
    });

    for production in prod_str {
        let it2 = production.split_ascii_whitespace();
        let mut idx: usize = 0;
        let mut ret: usize = usize::MAX;
        let mut rhs: Vec<usize> = vec![];

        let mut ok: bool = false;

        for c in it2 {
            if c == "::=" {
                if idx != 1 {
                    println!("\"::=\" must be at the second place.");
                    return Err(std::io::ErrorKind::InvalidInput.into());
                }
                ok = true;
            }

            if idx == 0 {
                ret = *map.get(c).unwrap();
                let val = related_lhs_prod_id.get_mut(ret).unwrap();
                val.insert(ps.len());
            } else if idx > 1 {
                rhs.push(*map.get(c).unwrap());
            }

            idx = idx + 1;
        }

        if !ok {
            println!("Production needs \"::=\" for indecating its structure.");
            return Err(std::io::ErrorKind::InvalidInput.into());
        }

        if idx <= 2 {
            println!("Production right hand side symbol sequence can not be empty.");
            return Err(std::io::ErrorKind::InvalidInput.into());
        }

        ps.push(Production { ret: ret, rhs: rhs });
    }

    if !map.contains_key(&start) {
        println!("The grammer must has a production with start symbol at left hand size.");
        return Err(std::io::ErrorKind::InvalidInput.into());
    }

    Ok(Grammer {
        sym: sym,
        prod: ps,
        nul: map[&nul],
        start: map[&start],
        sym_map: map,
        term_vis: term_vis,
        related_lhs_prod_id: related_lhs_prod_id,

        first: vec![],
        goto: HashMap::new(),
        lr0_kernels: BTreeSet::new(),
        lalr1_all_lookahead: vec![],
        lalr1_lookahead_set_to_id: BTreeMap::new(),
        lalr1_lookahead_map: BTreeMap::new(),
        lalr1_transmit_map: BTreeMap::new(),
    })
}
